#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   nmcli.py
@Time    :   2024-08-21 10:32
@Author  :   坐公交也用券
@Version :   1.0
@Contact :   faith01238@hotmail.com
@Homepage : https://liumou.site
@Desc    :   通过nmcli批量设置静态IP
"""
import logging
from argparse import ArgumentParser
from os import system
from subprocess import getstatusoutput
from sys import exit, stdout
from time import sleep

# 创建一个自定义的日志记录器
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# 创建一个自定义的日志处理器，设置其输出格式
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d | %(funcName)s | %(message)s')
handler = logging.StreamHandler(stdout)
handler.setFormatter(formatter)
logger.addHandler(handler)


def get_default_gw():
	"""
	获取默认网关
	:return:
	"""
	gw = getstatusoutput('ip route | grep default')
	if gw[0] == 0:
		gw = gw[1].split(' ')[2]
		return gw
	else:
		return None


class Nmcli:
	def __init__(self, arg: dict):
		"""

		:param arg: 参数字典
		"""
		self.gateway = arg['gateway']
		self.netmask = arg['netmask']
		self.dns1 = arg['dns1']
		self.dns2 = arg['dns2']
		self.sub = arg['sub']
		# 需要操作的网卡
		self.iface = ''
		# 需要操作的网卡列表
		self.iface_list = []
		self.iface_dict = {}
		self.iface_connect = {}
		self.iface_ip_dict = {}  # 记录网卡对应的ip
		self.iface_connect_name_uuid = {}
		self.iface_connect_uuid_name = {}
		self.iface_connect_name_name = {}
		self.iface_connect_uuid_list = []
		self.iface_id_name = {}
		self.logger = logger

	def check_ip(self):
		"""
		检查IP是否符合V4标准
		:return:
		"""
		gws = self.gateway.split('.')
		if len(gws) != 4:
			self.logger.error('网关IP格式错误')
			exit(1)
		else:
			for i in gws:
				if int(i) < 0 or int(i) > 255:
					self.logger.error('网关IP格式错误')
					exit(1)
		# 检查网段
		subs = self.sub.split('.')
		if len(subs) != 3:
			self.logger.error(f'网段格式错误: 网段分段个数错误,标准为3,当前：{len(subs)}')
			exit(1)
		else:
			for i in subs:
				if int(i) < 0 or int(i) > 255:
					self.logger.debug(i)
					self.logger.error('网段格式错误')
					exit(1)
		# # 检查掩码
		# nms = self.netmask.split('.')
		# if len(nms) != 4:
		# 	self.logger.error('掩码格式错误')
		# 	exit(1)
		# else:
		# 	for i in nms:
		# 		if int(i) < 0 or int(i) > 255:
		# 			self.logger.error('掩码格式错误')
		# 			exit(1)
		self.logger.info('IP配置信息检测成功')

	def get_connect_uuid(self):
		"""
		获取连接的uuid
		:return:
		"""
		cmd_ = "nmcli connection  | awk '{print $2,$4}' | grep -v UUID | grep -v docker | grep -v lo"
		res = getstatusoutput(cmd_)
		if res[0] == 0:
			for i in res[1].split('\n'):
				uuid = i.split(' ')[0]
				if len(uuid) != 36:
					self.logger.warning('UUID格式错误')
					continue
				self.iface_connect_uuid_list.append(uuid)

				eth_ = i.split(' ')[1]
				if len(eth_) > 3:
					self.logger.warning('eth名称格式错误')
					continue
				logger.debug(f'uuid: {uuid} - eth: {eth_}')
				self.iface_connect_name_uuid[eth_] = uuid
				self.iface_connect_uuid_name[uuid] = eth_
		# 获取网卡对应连接名称
		cmd_ = "nmcli connection  | awk '{print $1,$4}' | grep -v UUID | grep -v docker | grep -v lo"
		res = getstatusoutput(cmd_)
		if res[0] == 0:
			for i in res[1].split('\n'):
				name_ = i.split(' ')[0]
				eth_ = i.split(' ')[1]
				self.iface_connect_name_name[eth_] = name_

	def init_get_iface(self):
		"""
		获取网卡列表
		:return:
		"""
		index_ = 1
		self.logger.info('获取网卡列表')
		eth_ = getstatusoutput('nmcli device show')
		if eth_[0] == 0:
			eth = eth_[1].split('\n')
			for i in eth:
				if 'GENERAL.DEVICE:' in i:
					# 排除 lo docker
					if 'lo' not in i and 'docker' not in i:
						name = str(i.split(':')[1]).strip()
						self.logger.debug(name)
						self.iface_list.append(name)
						self.iface_dict[index_] = name
						self.iface_id_name[index_] = name
						index_ += 1
		else:
			print(eth_[1])
			self.logger.error('获取网卡列表失败')
			exit(1)
		self.logger.info(f'网卡列表: {self.iface_list}')

	def delete_old_connect(self, eth, name):
		"""
		删除旧的连接
		:param eth: 需要删除的网卡
		:param name: 需要设置的连接名称
		:return:
		"""
		self.logger.info('删除旧的连接')
		for uuid, eth_name in self.iface_connect_uuid_name.items():
			if eth == eth_name:
				self.logger.info(f'删除旧的连接: {eth}')
				cmd = f'nmcli connection delete {uuid}'
				if system(cmd) != 0:
					print(cmd)
					self.logger.error(f'删除旧的连接失败: {eth}')
					exit(1)

	def _get_connect_uuid(self):
		"""
		获取连接的uuid
		:return:
		"""
		cmd = "nmcli connection show %s | grep uuid | awk '{print $2}'" % self.iface
		res = getstatusoutput(cmd)
		if res[0] == 0:
			return res[1]
		else:
			print(cmd)
			self.logger.error(f'获取连接的uuid失败: {res[1]}')
			exit(1)

	def set_ip(self, ip):
		"""
		设置IP
		:param ip:
		:return:
		"""
		self.delete_old_connect(self.iface, self.iface)
		self.logger.info(f'设置IP: {ip} 网卡: {self.iface}')
		add_ = f"nmcli connection add type ethernet con-name {self.iface}  ifname {self.iface}"
		if system(add_) != 0:
			print(add_)
			self.logger.error('添加连接失败(Adding connection failed)')
			exit(1)

		uuid = self._get_connect_uuid()
		cmd = f'nmcli connection modify {uuid} ipv4.addresses {ip}/{self.netmask} ipv4.gateway {self.gateway} ipv4.method manual'
		if system(cmd) != 0:
			print(cmd)
			self.logger.error('设置IP失败(Failed to set IP address)')
			exit(1)
		# 增加DNS
		cmd = f'nmcli connection modify {uuid} ipv4.dns {self.dns1}'
		if system(cmd) != 0:
			print(cmd)
			self.logger.error(f'设置DNS失败: {self.dns1}')
			exit(1)
		# 设置第二个DNS
		cmd = f'nmcli connection modify {uuid} +ipv4.dns {self.dns2}'
		if system(cmd) != 0:
			print(cmd)
			self.logger.error(f'设置DNS失败: {self.dns2}')
			exit(1)
		# 启用连接
		cmd = f'nmcli connection up {uuid}'
		if system(cmd) != 0:
			print(cmd)
			self.logger.error(f'启用连接失败: {self.iface}')
			exit(1)
		self.logger.info(f'设置IP成功: {ip} 网卡: {self.iface}')
		system('nmcli connection')

	def mode2(self):
		"""

		:return:
		"""
		for k, v in self.iface_dict.items():
			self.logger.info(f'{k}. {v}')

	def ping(self, ip):
		"""

		:param ip:
		:return:
		"""
		if system(f'ping -c 3 {ip}') == 0:
			return True
		return False

	def mode1(self):
		"""
		单网卡模式
		:return:
		"""
		# 获取需要设置的IP,例如：100
		while True:
			try:
				print(
					"Please enter the last digit of the IP address, there is no need to enter the complete IP address, for example, 100")
				ip = int(input(f'请输入IP尾数,不需要输入完整IP,例如100, {self.sub}.: \n'))
				if ip < 0 or ip > 254:
					self.logger.error('IP范围错误(IP range error)')
					continue
				if self.ping(f'{self.sub}.{ip}'):
					logger.warning(f'{ip} IP地址被占用(The IP address is occupied)')
					continue
				self.set_ip(f'{self.sub}.{ip}')
				break
			except ValueError:
				self.logger.error(f'IP格式错误: {ValueError}')
				continue

	def mode1_get(self):
		while True:
			print('Please enter the corresponding serial number of the network card')
			print('请输入网卡对应的序号')
			for i, v in self.iface_dict.items():
				print(f'{i}. {v} - {self.iface_ip_dict[v]}')

			eth_id = input('Please enter: \n')
			try:
				eth_id = int(eth_id)
			except ValueError:
				self.logger.error(f'序号格式错误: {ValueError}')
			if eth_id not in self.iface_dict:
				logger.error(f'序号不存在 (The serial number does not exist): {eth_id}')
				continue

			self.iface = self.iface_dict[eth_id]
			if self.iface in self.iface_list:
				self.logger.info(f'使用这个网卡进行配置 (Use this network card for configuration): {self.iface}')
				return True
			else:
				self.logger.error(f'网卡不存在 (The network card does not exist): {self.iface}')
			return False

	def delete_all_connect(self):
		"""
		删除所有连接
		:return:
		"""
		self.logger.info('删除所有连接')
		print(self.iface_connect_uuid_list)
		for uuid in self.iface_connect_uuid_list:
			if len(uuid) < 10:
				continue
			cmd_ = f'nmcli connection delete {uuid}'
			if system(cmd_) != 0:
				self.logger.error(f'删除连接失败: {uuid}')
				print(cmd_)
				exit(1)
			else:
				self.logger.info(f'删除连接成功: {uuid}')

	def restart_nmcli(self):
		self.logger.info('重启NetworkManager服务')
		cmd_ = 'systemctl restart NetworkManager'
		if system(cmd_) != 0:
			self.logger.error('重启服务失败')
			print(cmd_)
			exit(1)
		else:
			self.logger.info('重启服务成功')

	def close_ipv6(self):
		self.logger.info('关闭所有网卡的IPV6')
		for eth in self.iface_list:
			cmd_ = f'nmcli connection modify {eth} ipv6.method disabled'
			if system(cmd_) != 0:
				self.logger.error(f'关闭网卡的IPV6失败: {eth}')
				print(cmd_)
				exit(1)
			else:
				self.logger.info(f'关闭网卡的IPV6成功: {eth}')
		# 重启服务
		self.restart_nmcli()

	def open_ipv6(self):
		self.logger.info('开启所有网卡的IPV6')
		for eth in self.iface_list:
			cmd_ = f'nmcli connection modify {eth} ipv6.method auto'
			if system(cmd_) != 0:
				self.logger.error(f'开启网卡的IPV6失败: {eth}')
				print(cmd_)
				exit(1)
			else:
				self.logger.info(f'开启网卡的IPV6成功: {eth}')
		self.restart_nmcli()

	def select_mode(self):
		self.logger.info('开始选择操作模式')
		if len(self.iface_list) == 1:
			self.iface = self.iface_list[0]
			self.logger.info(f'网卡名称: {self.iface}')
			self.mode1()
			return True
		while True:
			self.init_get_eth_ip_dict()
			print('请选择操作模式')
			print('[1] 单网卡模式(Single network card mode)')
			print('[2] 多网卡模式(Multi network card mode)')
			print('[3] 删除所有连接(Delete all connections)')
			print('[4] 关闭所有网卡的IPV6(Turn off IPV6 for all network cards)')
			print('[5] 开启所有网卡的IPV6(Turn on IPV6 for all network cards)')
			print('[6] 开启指定网卡的IPV6(Turn on IPV6 for the specified network card)')
			print('[q] 退出(quit)')
			print('请输入需要选择的序号')
			mode = input('Please enter the serial number to be selected: \n')
			if mode == '1':
				if self.mode1_get():
					self.mode1()
				else:
					continue
			elif mode == '2':
				if self.mode2():
					break
			elif mode == '3':
				self.delete_all_connect()
			elif mode == '4':
				self.close_ipv6()
			elif mode == '5':
				self.open_ipv6()
			elif mode == '6':
				logger.warning('待完善...')
			elif mode == 'q':
				exit(0)
			else:
				self.logger.error('请输入正确的模式(Please enter the correct mode)')

	def init_get_eth_ip_dict(self):
		"""
		初始化网卡IP字典
		:return:
		"""
		if len(self.iface_list) == 0:
			self.logger.error('没有网卡(No network card)')
			exit(1)
		# 通过网卡名称获取网卡IP
		for eth in self.iface_list:
			cmd = f'ip addr show {eth} | grep inet | awk \'{{print $2}}\''
			res = getstatusoutput(cmd)
			if res[0] == 0:
				self.iface_ip_dict[eth] = res[1].split('/')[0]
			else:
				self.logger.error(f'获取网卡IP失败: {eth}')
				print(cmd)
				exit(1)

	def start(self):
		"""
		开始操作
		:return:
		"""
		self.logger.info('开始操作(Start operation)')
		self.check_ip()
		self.init_get_iface()
		self.init_get_eth_ip_dict()
		self.get_connect_uuid()
		self.select_mode()


if __name__ == '__main__':
	ver = '1.4'
	sub_ = '192.168.1'
	gw_ = '192.168.1.1'
	gw = get_default_gw()
	if gw is None:
		logger.warning('获取默认网关失败')
		logger.warning("Failed to obtain default gateway")
	else:
		sub_ = f"{gw.split('.')[0]}.{gw.split('.')[1]}.{gw.split('.')[2]}"
		gw_ = gw
	parser = ArgumentParser(description=f'通过nmcli批量设置静态IP: {ver} 作者: 坐公交也用券')
	parser.add_argument('-s', '--sub', help=f'网段,默认：{sub_}', required=False, default=sub_)
	parser.add_argument('-g', '--gateway', help=f'网关,默认: {gw_}', required=False, default=gw_)
	parser.add_argument('-n', '--netmask', help='子网掩码,默认: 24', required=False, default=24, type=int)
	parser.add_argument('-d1', '--dns1', help='DNS', required=False, default='119.29.29.29')
	parser.add_argument('-d2', '--dns2', help='DNS', required=False, default='223.5.5.5')
	args = parser.parse_args()
	args_dict = vars(args)
	print(args_dict)
	logger.info(f"版本: {ver}")
	ns = Nmcli(args_dict)
	ns.start()
